So, as I mentioned on Friday I'm running into some difficulties with applying axis tags to the computations for two-point fluxes (link) / volume flux differencing (link). In the code in the first link, the incoming state_ll and state_rr are per-element row and column vectors (i.e. (nelements, 1, ndofs) and (nelements, ndofs, 1)). They're added together to produce an (nelements, ndofs, ndofs) array via broadcasting, which is then used in subsequent operations. When both axes are tagged with DiscretizationDOFAxisTag, I get the following error:

  File "meshmode/array_context.py", line 904, in fuse_same_discretization_entity_loops
    knl = _fuse_loops_over_a_discr_entity(knl, DiscretizationDOFAxisTag,
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "meshmode/array_context.py", line 876, in _fuse_loops_over_a_discr_entity
    lp.get_kennedy_unweighted_fusion_candidates(
  File "loopy/transform/loop_fusion.py", line 762, in get_kennedy_unweighted_fusion_candidates
    raise LoopyError(f"'{iname}' and '{conflict_iname}' "
loopy.diagnostic.LoopyError: 'cse_294_dim2' and 'cse_294_dim1' cannot fused be fused as they can be nested within one another.

Based on our discussion on Friday it sounds like the fusion pipeline isn't yet smart enough to avoid attempting to fuse loops over the first DOF axis with loops over the second. Previously I had attempted to fix this by creating a new DOF axis tag and applying that to the second DOF axis of the arrays, but this appears to just end up producing duplicate tags once the axis is inferred as a DiscretizationDOFAxisTag from context with the other arrays involved in the computation.

I'm not sure what the next step should be. Question @inducer: Is there a reason that the failure to fuse needs to be an error here vs. just a warning? Would changing that be one way of resolving this, or does that not help? Anything else I can try?

@inducer Any thoughts on this?

For the record (and as discussed) this comes from the current transform strategy not into account that there could be more than one axis corresponding to DOFs in an array. Also, it seems to want to materialize the 2D DOF differencing array, which is almost certainly not what we want.

Also, it seems to want to materialize the 2D DOF differencing array, which is almost certainly not what we want.

I'm skeptical if that's true with the newly proposed BatchedEinsumArrayContext as well. The arguments of an einsum node are tagged with ImplSubstitution in that version.